#!/usr/bin/env python3
"""
屏幕截图工具
"""
import os
import sys
import time
import asyncio
import tempfile
from concurrent.futures import ThreadPoolExecutor

from PySide6.QtWidgets import QWidget, QApplication, QLabel, QRubberBand
from PySide6.QtGui import QPainter, QColor, QPen, QCursor, QPixmap, QKeyEvent
from PySide6.QtCore import Qt, QRect, QSize, Signal, QPoint

from src.utils.style_manager import StyleManager

# 添加项目根目录到路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../..")))

class ScreenshotSelector(QWidget):
    """截图选择器，允许用户选择屏幕区域"""
    region_selected = Signal(QRect)
    canceled = Signal()
    
    def __init__(self):
        super().__init__()
        self.setWindowFlags(Qt.FramelessWindowHint | Qt.WindowStaysOnTopHint)
        self.setWindowState(Qt.WindowFullScreen)  # 确保是全屏窗口
        self.setCursor(Qt.CrossCursor)
        
        # 获取屏幕信息
        self.screen = QApplication.primaryScreen()
        self.device_pixel_ratio = self.screen.devicePixelRatio()
        self.screen_rect = self.screen.geometry()
        self.screen_width = self.screen_rect.width()
        self.screen_height = self.screen_rect.height()
        
        # 获取屏幕截图 - 使用原始尺寸 
        self.original_pixmap = self.screen.grabWindow(0)
        self.original_width = self.original_pixmap.width()
        self.original_height = self.original_pixmap.height()
        
        # 打印调试信息
        StyleManager.debug_print(f"原始图像尺寸: {self.original_width}x{self.original_height}")
        StyleManager.debug_print(f"屏幕尺寸: {self.screen_width}x{self.screen_height}")
        StyleManager.debug_print(f"设备像素比: {self.device_pixel_ratio}")
        
        # 创建用于显示的截图 - 确保填满整个屏幕
        self.screenshot = QPixmap(self.screen_width, self.screen_height)
        self.screenshot.fill(Qt.transparent)
        
        painter = QPainter(self.screenshot)
        
        # 根据设备像素比调整截图大小，确保填满整个屏幕
        if self.original_width != self.screen_width or self.original_height != self.screen_height:
            StyleManager.debug_print(f"调整截图大小以匹配屏幕: {self.screen_width}x{self.screen_height}")
            
            # 计算绘制位置，使图像居中
            x = 0
            y = 0
            
            # 绘制原始图像，拉伸以匹配屏幕大小
            target_rect = QRect(x, y, self.screen_width, self.screen_height)
            source_rect = QRect(0, 0, self.original_width, self.original_height)
            painter.drawPixmap(target_rect, self.original_pixmap, source_rect)
        else:
            # 直接绘制原始图像
            painter.drawPixmap(0, 0, self.original_pixmap)
        
        painter.end()
        
        # 初始化选区相关变量
        self.origin = QPoint()
        self.current_rect = QRect()
        
        # 设置帮助文本
        self.help_label = QLabel(self)
        self.help_label.setText("拖动鼠标选择区域，按Esc取消，Enter确认")
        self.help_label.setStyleSheet("""
            QLabel {
                color: white;
                font-size: 16px;
                background-color: rgba(0, 0, 0, 150);
                padding: 8px 12px;
                border-radius: 4px;
            }
        """)
        self.help_label.move(10, 10)
        self.help_label.adjustSize()
        
        # 尺寸提示标签
        self.size_label = QLabel(self)
        self.size_label.setStyleSheet("""
            QLabel {
                color: white;
                font-size: 14px;
                background-color: rgba(0, 0, 0, 150);
                padding: 4px 8px;
                border-radius: 4px;
            }
        """)
        self.size_label.hide()
    
    def showEvent(self, event):
        """显示事件，设置截图选择器位置和大小"""
        super().showEvent(event)
        
        # 设置窗口大小为全屏
        self.setGeometry(self.screen_rect)
        
        # 设置窗口为活动窗口，确保它位于最前面
        self.activateWindow()
        self.raise_()
        
        StyleManager.debug_print(f"显示截图选择器，当前大小: {self.width()}x{self.height()}")
    
    def paintEvent(self, event):
        """绘制事件，显示截图背景和选区"""
        painter = QPainter(self)
        painter.setRenderHint(QPainter.Antialiasing)
        
        # 绘制背景截图
        painter.drawPixmap(0, 0, self.screenshot)
        
        if not self.current_rect.isNull():
            # 绘制半透明遮罩
            mask = QColor(0, 0, 0, 100)
            painter.setBrush(mask)
            painter.setPen(Qt.NoPen)
            
            # 绘制四个遮罩区域
            painter.drawRect(0, 0, self.width(), self.current_rect.top())  # 上方
            painter.drawRect(0, self.current_rect.bottom(), self.width(), self.height() - self.current_rect.bottom())  # 下方
            painter.drawRect(0, self.current_rect.top(), self.current_rect.left(), self.current_rect.height())  # 左侧
            painter.drawRect(self.current_rect.right(), self.current_rect.top(), 
                           self.width() - self.current_rect.right(), self.current_rect.height())  # 右侧
            
            # 绘制选区边框
            pen = QPen(QColor("#409EFF"), 2)
            painter.setPen(pen)
            painter.setBrush(Qt.NoBrush)
            painter.drawRect(self.current_rect)
            
            # 绘制选区四角
            corner_size = 10
            corner_width = 2
            corner_color = QColor("#409EFF")
            painter.setPen(QPen(corner_color, corner_width))
            
            # 左上角
            painter.drawLine(self.current_rect.topLeft().x(), self.current_rect.topLeft().y(),
                           self.current_rect.topLeft().x() + corner_size, self.current_rect.topLeft().y())
            painter.drawLine(self.current_rect.topLeft().x(), self.current_rect.topLeft().y(),
                           self.current_rect.topLeft().x(), self.current_rect.topLeft().y() + corner_size)
            
            # 右上角
            painter.drawLine(self.current_rect.topRight().x(), self.current_rect.topRight().y(),
                           self.current_rect.topRight().x() - corner_size, self.current_rect.topRight().y())
            painter.drawLine(self.current_rect.topRight().x(), self.current_rect.topRight().y(),
                           self.current_rect.topRight().x(), self.current_rect.topRight().y() + corner_size)
            
            # 左下角
            painter.drawLine(self.current_rect.bottomLeft().x(), self.current_rect.bottomLeft().y(),
                           self.current_rect.bottomLeft().x() + corner_size, self.current_rect.bottomLeft().y())
            painter.drawLine(self.current_rect.bottomLeft().x(), self.current_rect.bottomLeft().y(),
                           self.current_rect.bottomLeft().x(), self.current_rect.bottomLeft().y() - corner_size)
            
            # 右下角
            painter.drawLine(self.current_rect.bottomRight().x(), self.current_rect.bottomRight().y(),
                           self.current_rect.bottomRight().x() - corner_size, self.current_rect.bottomRight().y())
            painter.drawLine(self.current_rect.bottomRight().x(), self.current_rect.bottomRight().y(),
                           self.current_rect.bottomRight().x(), self.current_rect.bottomRight().y() - corner_size)
            
            # 更新尺寸提示
            size_text = f"{self.current_rect.width()} × {self.current_rect.height()}"
            self.size_label.setText(size_text)
            self.size_label.adjustSize()
            
            # 调整尺寸提示位置
            label_x = self.current_rect.right() - self.size_label.width()
            label_y = self.current_rect.bottom() + 5
            if label_y + self.size_label.height() > self.height():
                label_y = self.current_rect.top() - self.size_label.height() - 5
            self.size_label.move(label_x, label_y)
            self.size_label.show()
        else:
            self.size_label.hide()
    
    def mousePressEvent(self, event):
        """鼠标按下事件，开始选择区域"""
        if event.button() == Qt.LeftButton:
            self.origin = event.pos()
            self.current_rect = QRect()
            self.update()
    
    def mouseMoveEvent(self, event):
        """鼠标移动事件，更新选择区域"""
        if not self.origin.isNull():
            self.current_rect = QRect(self.origin, event.pos()).normalized()
            self.update()
    
    def mouseReleaseEvent(self, event):
        """鼠标释放事件，完成选择区域"""
        if event.button() == Qt.LeftButton:
            self.current_rect = QRect(self.origin, event.pos()).normalized()
            
            # 只有当选择区域足够大时才触发选择完成
            if self.current_rect.width() > 10 and self.current_rect.height() > 10:
                # 发送原始的选择区域
                self.region_selected.emit(self.current_rect)
                # 自动关闭窗口，无需用户再按Enter
                self.close()
            self.update()
    
    def keyPressEvent(self, event):
        """键盘事件，处理ESC取消和Enter确认"""
        if event.key() == Qt.Key_Escape:
            self.canceled.emit()
            self.close()
        elif event.key() == Qt.Key_Return or event.key() == Qt.Key_Enter:
            if not self.current_rect.isNull() and self.current_rect.width() > 10 and self.current_rect.height() > 10:
                # 发送原始的选择区域
                self.region_selected.emit(self.current_rect)
            else:
                self.canceled.emit()
            self.close()

class ScreenshotUtil:
    # 线程池执行器，用于执行耗时操作
    _executor = ThreadPoolExecutor(max_workers=2)
    
    @staticmethod
    async def capture_screen_async(region=True, temp_dir=None):
        """异步截取屏幕，可选区域截图或全屏截图
        
        Args:
            region: 是否使用区域截图，默认为True
            temp_dir: 临时文件保存目录，默认使用系统临时目录
            
        Returns:
            元组 (temp_file, pixmap, rect): 临时文件路径、QPixmap对象和区域位置
            对于全屏截图，rect将是(0,0)开始的整个屏幕区域
        """
        if region:
            # 区域截图
            return await ScreenshotUtil._capture_region_async(temp_dir)
        else:
            # 全屏截图
            temp_file, pixmap = await asyncio.get_event_loop().run_in_executor(
                ScreenshotUtil._executor,
                ScreenshotUtil._capture_full_screen,
                temp_dir
            )
            return temp_file, pixmap, None
            
    @staticmethod
    async def capture_region_async(x, y, width, height, output_path=None):
        """异步捕获指定区域的屏幕
        
        Args:
            x: 区域左上角的x坐标
            y: 区域左上角的y坐标
            width: 区域宽度
            height: 区域高度
            output_path: 输出文件路径，如果为None则使用临时文件
            
        Returns:
            成功返回True，失败返回False
        """
        try:
            StyleManager.debug_print(f"捕获区域截图: ({x}, {y}, {width}, {height})")
            
            # 确保宽度和高度为正值
            if width <= 0 or height <= 0:
                StyleManager.debug_print("无效的宽度或高度")
                return False
                
            # 如果未指定输出路径，创建临时文件
            if not output_path:
                # 确保临时目录存在
                temp_dir = os.path.join(tempfile.gettempdir(), "PySide6x")
                os.makedirs(temp_dir, exist_ok=True)
                
                # 创建临时文件
                output_path = os.path.join(temp_dir, f"region_screenshot_{int(time.time())}.png")
            
            # 使用线程池执行截图操作，避免阻塞UI线程
            def capture_region():
                try:
                    screen = QApplication.primaryScreen()
                    rect = QRect(x, y, width, height)
                    pixmap = screen.grabWindow(0, x, y, width, height)
                    
                    if pixmap.isNull():
                        StyleManager.debug_print("截图失败，返回空白图像")
                        return False
                        
                    # 保存截图
                    if pixmap.save(output_path):
                        StyleManager.debug_print(f"截图已保存到: {output_path}")
                        return True
                    else:
                        StyleManager.debug_print(f"无法保存截图到: {output_path}")
                        return False
                except Exception as e:
                    StyleManager.debug_print(f"截图过程中出错: {str(e)}")
                    return False
            
            # 异步执行截图
            result = await asyncio.get_event_loop().run_in_executor(
                ScreenshotUtil._executor,
                capture_region
            )
            
            return result
            
        except Exception as e:
            StyleManager.debug_print(f"捕获区域截图时出错: {str(e)}")
            return False

    @staticmethod
    async def _capture_region_async(temp_dir=None):
        """异步区域截图实现"""
        # 创建future对象，以便从UI回调中获取结果
        rect_future = asyncio.Future()
        
        # 创建选择器
        selector = ScreenshotSelector()
        
        # 连接信号
        def on_region_selected(rect):
            if not rect_future.done():
                rect_future.set_result(rect)
        
        def on_canceled():
            if not rect_future.done():
                rect_future.set_result(None)
        
        selector.region_selected.connect(on_region_selected)
        selector.canceled.connect(on_canceled)
        
        # 显示选择器
        selector.show()
        
        # 等待用户选择
        try:
            rect = await rect_future
            if rect:
                # 计算选择区域在原始图像上的确切位置
                StyleManager.debug_print(f"屏幕上选择区域: x={rect.x()}, y={rect.y()}, w={rect.width()}, h={rect.height()}")
                
                # 如果屏幕尺寸与原始图像尺寸不一致，需要计算原始图像中的对应区域
                source_rect = rect
                if selector.original_width != selector.screen_width or selector.original_height != selector.screen_height:
                    # 计算比例
                    ratio_x = selector.original_width / selector.screen_width
                    ratio_y = selector.original_height / selector.screen_height
                    
                    # 计算在原始图像中的坐标和尺寸
                    source_x = int(rect.x() * ratio_x)
                    source_y = int(rect.y() * ratio_y)
                    source_width = int(rect.width() * ratio_x)
                    source_height = int(rect.height() * ratio_y)
                    
                    source_rect = QRect(source_x, source_y, source_width, source_height)
                    StyleManager.debug_print(f"原始图像上对应区域: x={source_x}, y={source_y}, w={source_width}, h={source_height}")
                
                # 确保区域不超出原始图像边界
                pixmap_rect = selector.original_pixmap.rect()
                bounded_rect = source_rect.intersected(pixmap_rect)
                
                # 如果选择区域不同，记录日志
                if bounded_rect != source_rect:
                    StyleManager.debug_print(f"选择区域超出范围，调整为: x={bounded_rect.x()}, y={bounded_rect.y()}, w={bounded_rect.width()}, h={bounded_rect.height()}")
                
                # 从原始图像中直接提取选择区域
                selected_pixmap = selector.original_pixmap.copy(bounded_rect)
                
                # 保存图片
                if temp_dir is None:
                    temp_dir = tempfile.gettempdir()
                os.makedirs(temp_dir, exist_ok=True)
                temp_file = os.path.join(temp_dir, "region_capture_" + str(int(time.time())) + ".png")
                
                # 使用PNG格式保存，保持完全的图像质量 - 无压缩
                selected_pixmap.save(temp_file, "PNG", -1)
                
                # 打印截图信息用于调试
                StyleManager.debug_print(f"最终截图尺寸: {selected_pixmap.width()}x{selected_pixmap.height()}")
                StyleManager.debug_print(f"保存截图到: {temp_file}")
                
                # 返回元组：(临时文件路径, 截图QPixmap, 屏幕上的截图区域QRect)
                return temp_file, selected_pixmap, rect
            return None, None, None
        except Exception as e:
            StyleManager.debug_print(f"区域截图错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return None, None, None

    @staticmethod
    def _capture_full_screen(temp_dir=None):
        """全屏截图
        
        Args:
            temp_dir: 临时文件保存目录，默认使用系统临时目录
            
        Returns:
            元组 (temp_file, pixmap): 临时文件路径和QPixmap对象
        """
        try:
            # 创建截图
            screen = QApplication.primaryScreen()
            pixmap = screen.grabWindow(0)
            
            # 保存图片
            if temp_dir is None:
                temp_dir = tempfile.gettempdir()
            os.makedirs(temp_dir, exist_ok=True)
            temp_file = os.path.join(temp_dir, "screen_capture_" + str(int(time.time())) + ".png")
            
            # 使用PNG格式保存，保持完全的图像质量 - 无压缩
            pixmap.save(temp_file, "PNG", -1)
            
            # 打印截图信息用于调试
            StyleManager.debug_print(f"全屏截图尺寸: {pixmap.width()}x{pixmap.height()}")
            StyleManager.debug_print(f"保存截图到: {temp_file}")
            
            return temp_file, pixmap
        except Exception as e:
            StyleManager.debug_print(f"全屏截图错误: {str(e)}")
            import traceback
            traceback.print_exc()
            return None, None 